package com.sunlake.spring.main.common.utils;

import com.sunlake.spring.main.config.JwtConfig;
import com.sunlake.spring.main.model.auth.po.UsersPO;
import io.jsonwebtoken.Claims;
import io.jsonwebtoken.Jwts;
import io.jsonwebtoken.security.Keys;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.data.redis.core.RedisTemplate;
import org.springframework.stereotype.Component;
import javax.annotation.PostConstruct;
import javax.crypto.SecretKey;
import java.security.PrivateKey;
import java.security.PublicKey;
import java.util.Date;
import java.util.HashMap;
import java.util.Map;
import java.util.concurrent.TimeUnit;

/**
 * JWT令牌工具类
 * 支持 HMAC(HS512) 与 RSA(RS512) 两种模式。根据 JwtConfig 中是否提供 RSA Key 自动选择。
 */
@Component
public class JwtTokenUtil {
    @Autowired
    private JwtConfig jwtConfig;

    @Autowired
    private RedisTemplate<String, Object> redisTemplate;

    // HMAC 访问/刷新密钥
    private SecretKey accessTokenKey;
    private SecretKey refreshTokenKey;

    // RSA 密钥对
    private PrivateKey rsaPrivateKey;
    private PublicKey rsaPublicKey;
    private boolean useRsa = false;

    private static final String CLAIM_KEY_USERNAME = "username";
    private static final String CLAIM_KEY_CREATED = "created";
    private static final String CLAIM_KEY_USER_ID = "userId";

    @PostConstruct
    public void init() {
        tryLoadRsaKeys();
        if (!useRsa) {
            String accessKeyString = jwtConfig.getSecret();
            this.accessTokenKey = ensureSecureKeySize(accessKeyString);
            String refreshKeyString = jwtConfig.getSecret() + "_refresh";
            this.refreshTokenKey = ensureSecureKeySize(refreshKeyString);
        }
    }

    private void tryLoadRsaKeys() {
        try {
            String privatePem = null;
            String publicPem = null;
            if (jwtConfig.getPrivateKeyPath() != null && !jwtConfig.getPrivateKeyPath().isBlank()) {
                privatePem = RsaKeyUtil.readAll(java.nio.file.Path.of(jwtConfig.getPrivateKeyPath()));
            } else if (jwtConfig.getPrivateKeyPem() != null && !jwtConfig.getPrivateKeyPem().isBlank()) {
                privatePem = jwtConfig.getPrivateKeyPem();
            }
            if (jwtConfig.getPublicKeyPath() != null && !jwtConfig.getPublicKeyPath().isBlank()) {
                publicPem = RsaKeyUtil.readAll(java.nio.file.Path.of(jwtConfig.getPublicKeyPath()));
            } else if (jwtConfig.getPublicKeyPem() != null && !jwtConfig.getPublicKeyPem().isBlank()) {
                publicPem = jwtConfig.getPublicKeyPem();
            }
            if (privatePem != null && publicPem != null) {
                rsaPrivateKey = RsaKeyUtil.parsePrivateKeyPem(privatePem);
                rsaPublicKey = RsaKeyUtil.parsePublicKeyPem(publicPem);
                useRsa = true;
            }
        } catch (Exception ignore) {
            useRsa = false;
        }
    }

    private SecretKey ensureSecureKeySize(String keyString) {
        if (keyString.getBytes().length < 64) {
            byte[] keyBytes = new byte[64];
            byte[] originalBytes = keyString.getBytes();
            System.arraycopy(originalBytes, 0, keyBytes, 0, Math.min(originalBytes.length, keyBytes.length));
            return Keys.hmacShaKeyFor(keyBytes);
        }
        return Keys.hmacShaKeyFor(keyString.getBytes());
    }

    // ================= 生成访问令牌 =================
    public String generateToken(Long userId, String username) {
        Map<String, Object> claims = baseClaims(userId, username);
        Date now = new Date();
        Date expiration = new Date(now.getTime() + jwtConfig.getExpireTime() * 1000);
        String token = buildJwt(claims, now, expiration, true);
        redisTemplate.opsForValue().set("token:" + username, token, jwtConfig.getExpireTime(), TimeUnit.SECONDS);
        return token;
    }

    public String generateToken(UsersPO userDetails) {
        Map<String, Object> claims = baseClaims(Long.valueOf(userDetails.getId()), userDetails.getUsername());
        if (userDetails.getRole() != null) {
            claims.put("role", userDetails.getRole().getCode());
        }
        Date now = new Date();
        Date expiration = new Date(now.getTime() + jwtConfig.getExpireTime() * 1000);
        String token = buildJwt(claims, now, expiration, true);
        redisTemplate.opsForValue().set("token:" + userDetails.getUsername(), token, jwtConfig.getExpireTime(), TimeUnit.SECONDS);
        return token;
    }

    // ================= 生成刷新令牌 =================
    public String generateRefreshToken(Long userId, String username) {
        Map<String, Object> claims = new HashMap<>();
        claims.put(CLAIM_KEY_USERNAME, username);
        claims.put(CLAIM_KEY_USER_ID, userId);
        Date now = new Date();
        Date expiration = new Date(now.getTime() + jwtConfig.getRefreshExpireTime() * 1000);
        String token = buildJwt(claims, now, expiration, false);
        redisTemplate.opsForValue().set("refresh_token:" + username, token, jwtConfig.getRefreshExpireTime(), TimeUnit.SECONDS);
        return token;
    }

    private Map<String, Object> baseClaims(Long userId, String username) {
        Map<String, Object> claims = new HashMap<>();
        claims.put(CLAIM_KEY_USERNAME, username);
        claims.put(CLAIM_KEY_CREATED, new Date());
        claims.put(CLAIM_KEY_USER_ID, userId);
        return claims;
    }

    private String buildJwt(Map<String, Object> claims, Date issuedAt, Date expiration, boolean access) {
        // 使用 JJWT 0.13+ 的新式流式 API（claims/issuedAt/expiration），不要强制转换为不存在的 HeaderBuilder
        io.jsonwebtoken.JwtBuilder builder = Jwts.builder()
                .claims(claims)
                .issuedAt(issuedAt)
                .expiration(expiration);
        if (useRsa) {
            builder.signWith(rsaPrivateKey, Jwts.SIG.RS512);
        } else {
            builder.signWith(access ? accessTokenKey : refreshTokenKey, Jwts.SIG.HS512);
        }
        return builder.compact();
    }

    // ================= Claims 解析 =================
    public Claims getClaimsFromToken(String token) {
        try {
            if (useRsa) {
                return Jwts.parser().verifyWith(rsaPublicKey).build().parseSignedClaims(token).getPayload();
            } else {
                return Jwts.parser().verifyWith(accessTokenKey).build().parseSignedClaims(token).getPayload();
            }
        } catch (Exception e) {
            return null;
        }
    }

    public Claims getClaimsFromRefreshToken(String token) {
        try {
            if (useRsa) {
                return Jwts.parser().verifyWith(rsaPublicKey).build().parseSignedClaims(token).getPayload();
            } else {
                return Jwts.parser().verifyWith(refreshTokenKey).build().parseSignedClaims(token).getPayload();
            }
        } catch (Exception e) {
            return null;
        }
    }

    public String getUserRoleFromToken(String token) {
        Claims claims = getClaimsFromToken(token);
        return claims == null ? null : claims.get("role", String.class);
    }

    public String getUsernameFromToken(String token) {
        Claims claims = getClaimsFromToken(token);
        return claims == null ? null : claims.get(CLAIM_KEY_USERNAME, String.class);
    }

    public Long getUserIdFromToken(String token) {
        Claims claims = getClaimsFromToken(token);
        return claims == null ? null : claims.get(CLAIM_KEY_USER_ID, Long.class);
    }

    public boolean isTokenExpired(String token) {
        Claims claims = getClaimsFromToken(token);
        if (claims == null) {
            return true;
        }
        Date expiration = claims.getExpiration();
        return expiration.before(new Date());
    }

    public boolean validateToken(String token, String username, Long userId) {
        if (username == null || userId == null) {
            return false;
        }
        Long tokenUserId = getUserIdFromToken(token);
        String tokenUsername = getUsernameFromToken(token);
        if (tokenUserId == null || tokenUsername == null) {
            return false;
        }
        if (!username.equals(tokenUsername) || !userId.equals(tokenUserId) || isTokenExpired(token)) {
            return false;
        }
        String redisToken = (String) redisTemplate.opsForValue().get("token:" + username);
        return token.equals(redisToken);
    }

    public String refreshToken(String refreshToken) {
        Claims claims = getClaimsFromRefreshToken(refreshToken);
        if (claims == null) {
            return null;
        }
        Long userId = claims.get(CLAIM_KEY_USER_ID, Long.class);
        String username = claims.get(CLAIM_KEY_USERNAME, String.class);
        if (userId == null || username == null || isRefreshTokenExpired(refreshToken)) {
            return null;
        }
        String redisRefreshToken = (String) redisTemplate.opsForValue().get("refresh_token:" + username);
        if (!refreshToken.equals(redisRefreshToken)) {
            return null;
        }
        return generateToken(userId, username);
    }

    private boolean isRefreshTokenExpired(String refreshToken) {
        Claims claims = getClaimsFromRefreshToken(refreshToken);
        if (claims == null) {
            return true;
        }
        return claims.getExpiration().before(new Date());
    }

    public void logout(String username) {
        redisTemplate.delete("token:" + username);
        redisTemplate.delete("refresh_token:" + username);
    }
}